package com.bo.day20241121;

import java.util.*;

/**
 * Kruskal算法
 * 生成最小树
 * 思路：
 * 1. 将所有边按照权值从小到大排序
 * 2. 遍历所有边，判断是否构成环，若构成环则舍弃，若不构成环则加入最小树
 * 判断是否环，可以用并查集，如果两个节点不属于同一个集合，则说明不构成环，否则构成环
 *
 * @Author: gpb
 * @Date: 2024/11/27 15:41
 * @Description:
 */
public class Kruskal {
    public static class UnionFind {
        public HashMap<Node, Node> fatherMap;
        public HashMap<Node, Integer> sizeMap;

        public UnionFind () {
            fatherMap = new HashMap<>();
            sizeMap = new HashMap<>();
        }

        /**
         * 初始化并查集
         *
         * @param nodes
         */
        public void makeSets (Collection<Node> nodes) {
            fatherMap.clear();
            sizeMap.clear();
            for (Node node : nodes) {
                fatherMap.put(node, node);
                sizeMap.put(node, 1);
            }
        }

        public Node findFather (Node node) {
            Stack<Node> path = new Stack<>();
            while (node != fatherMap.get(node)) {
                node = fatherMap.get(node);
                path.push(node);
            }
            while (!path.isEmpty()) {
                Node pop = path.pop();
                fatherMap.put(pop, node);
            }
            return node;
        }

        public boolean isSameSet (Node a, Node b) {
            return findFather(a) == findFather(b);
        }

        public void union (Node a, Node b) {
            Node aF = findFather(a);
            Node bF = findFather(b);
            if (aF != bF) {
                int aSize = sizeMap.get(aF);
                int bSize = sizeMap.get(bF);
                if (aSize >= bSize) {
                    fatherMap.put(bF, aF);
                    sizeMap.put(aF, aSize + bSize);
                    sizeMap.remove(bF);
                } else {
                    fatherMap.put(aF, bF);
                    sizeMap.put(bF, aSize + bSize);
                    sizeMap.remove(aF);
                }
            }
        }
    }

    /**
     * 定义边的比较器, 按照权值从小到大排序
     */
    public static class EdgeComparator implements Comparator<Edge> {

        @Override
        public int compare (Edge o1, Edge o2) {
            return o1.width - o2.width;
        }
    }

    public static Set<Edge> kruskal (Graph graph) {
        UnionFind unionFind = new UnionFind();
        unionFind.makeSets(graph.nodes.values());
        // 小根堆, 按照权值从小到大排序
        PriorityQueue<Edge> priorityQueue = new PriorityQueue<>(new EdgeComparator());
        for (Edge edge : graph.edges) {
            priorityQueue.add(edge);
        }
        Set<Edge> result = new HashSet<>();
        while (!priorityQueue.isEmpty()) {
            Edge edge = priorityQueue.poll();
            // 判断是否构成环
            if (!unionFind.isSameSet(edge.from, edge.to)) {
                result.add(edge);
                // 合并
                unionFind.union(edge.from, edge.to);
            }
        }
        return result;
    }
}
